#include "data_loader.h"
#include <assert.h>

reverse_int(unsigned int i) {
    unsigned char ch1, ch2, ch3, ch4;
    ch1 = i & 255;
    ch2 = (i >> 8) & 255;
    ch3 = (i >> 16) & 255;
    ch4 = (i >> 24) & 255;
    return ((unsigned int)ch1 << 24) + ((unsigned int)ch2 << 16) + ((unsigned int)ch3 << 8) + ch4;
}

void read_images(const char* fileName, image** out, int *batch) {
    printf("==============read_images==========\n");
    printf("file name %s\n", fileName);
    FILE *fp = fopen(fileName, "r");
    assert(fp != NULL);
    unsigned int magic_number = 0;
    unsigned int number_of_images = 0;
    unsigned int n_rows = 0;
    unsigned int n_cols = 0;
    fread((char*)&magic_number, sizeof(magic_number), 1, fp);
    fread((char*)&number_of_images, sizeof(number_of_images), 1, fp);
    fread((char*)&n_rows, sizeof(n_rows), 1, fp);
    fread((char*)&n_cols, sizeof(n_cols), 1, fp);
    magic_number = reverse_int(magic_number);
    number_of_images = reverse_int(number_of_images);
    n_rows = reverse_int(n_rows);
    n_cols = reverse_int(n_cols);
    printf("magic_number = %d, number_of_images=%d, n_rows=%d, n_cols=%d\n", magic_number, number_of_images, n_rows, n_cols);
    int i, j;
    *batch = number_of_images;
    int size = n_cols * n_rows;
    *out = (image*)malloc(sizeof(image) * number_of_images);
    unsigned char *img = (unsigned char *)malloc(sizeof(unsigned char) * size);
    for(i = 0; i < number_of_images; i++)
    {
        image *normalized_img = (image*)malloc(sizeof(image));
        normalized_img->data = (float*)malloc(sizeof(float) * size);
        fread(img, sizeof(unsigned char), size, fp);
        for (j = 0; j < size; j++)
        {
            normalized_img->data[j] = ((float)img[j]) / 255 - 0.5;
        }
        normalized_img->h = n_rows;
        normalized_img->w = n_cols;
        (*out)[i] = *normalized_img;
    }
    free(img);
    fclose(fp);
    printf("read_images done!\n");
}

void read_labels(const char* fileName, int** out, int *batch) {
    printf("==============read_labels==========\n");
    printf("file name %s\n", fileName);
    FILE *fp = fopen(fileName, "r");
    assert(fp != NULL);
    unsigned int magic_number = 0;
    unsigned int number_of_images = 0;
    fread((char*)&magic_number, sizeof(magic_number), 1, fp);
    fread((char*)&number_of_images, sizeof(number_of_images), 1, fp);
    magic_number = reverse_int(magic_number);
    number_of_images = reverse_int(number_of_images);
    *batch = number_of_images;
    printf("magic_number = %d, number_of_images=%d\n", magic_number, number_of_images);
    int i;
    *out = (int*)malloc(sizeof(int) * number_of_images);
    for (i = 0; i < number_of_images; i++)
    {
        unsigned char *label = (unsigned char*)malloc(sizeof(unsigned char));
        fread(label, sizeof(unsigned char), 1, fp);
        (*out)[i] = (int)*label;
    }
    fclose(fp);
    printf("read_labels done!\n");
}